#include <string>
#include <vector>
#include <iostream>
#include <fstream>
#include <iomanip>
#include <algorithm>
#include "../fft_gpu.cu"
#include "../utils.cuh"
#include "../convfft.cu"
#include "convfft_gpu.cu"

#pragma once

void conv_back_filter_fft_2d(
    Real* input_X, int in_n, int in_c, int in_d, int in_h, int in_w,
    Real* output_Y,int out_n, int out_c, int out_d, int out_h, int out_w,
    int p_h, int p_w,
    int s_h, int s_w,
    int d_h, int d_w,
    Real* input_W, int k_n, int k_c, int k_d, int k_h, int k_w)
{

    Real* input_X_trans = (Real*) malloc(in_n * in_c * in_d * in_h * in_w * sizeof(Real));
    Real* output_Y_trans = (Real*) malloc(out_n * out_c * out_d * out_h * out_w * sizeof(Real));
    Real* input_W_trans = (Real*) malloc( k_n * k_c * k_d * k_h * k_w*sizeof(Real));

    transposeByBlock(input_X,input_X_trans, in_n, in_c, in_h, in_w);
    transposeByBlock(output_Y,output_Y_trans, out_n, out_c, out_h, out_w);
    rot180Degree(output_Y_trans,out_n, out_c, out_d, out_h,  out_w);
    convfft_2d(input_X_trans, in_c, in_n, in_d, in_h, in_w,
        output_Y_trans, out_c,out_n,  out_d, out_h,  out_w,
         p_h,  p_w,
         s_h,  s_w,
         d_h,  d_w,
         input_W_trans,  k_c, k_n,k_d, k_h, k_w);
    transposeByBlock( input_W_trans ,input_W,  k_c, k_n,k_h, k_w);
   
    rot180Degree(input_W,k_n, k_c, k_d, k_h, k_w);

        free(input_X_trans);
        free(output_Y_trans);
        free(input_W_trans);
}


void conv_back_data_fft_2d(
    Real* output_Y,int out_n, int out_c, int out_d, int out_h, int out_w,
    Real* input_W, int k_n, int k_c, int k_d, int k_h, int k_w,
    int p_h, int p_w,
    int s_h, int s_w,
    int d_h, int d_w,
    Real* input_X, int in_n, int in_c, int in_d, int in_h, int in_w
   )
{
    rot180Degree(input_W,  k_n , k_c , k_d , k_h , k_w);

    convfft_2d(output_Y, out_n, out_c, out_d, out_h,  out_w,
        input_W, k_n, k_c, k_d, k_h, k_w,
         p_h,  p_w,
         s_h,  s_w,
         d_h,  d_w,
         input_X, in_n, in_c, in_d, in_h, in_w);

}


void conv_back_data_fft_3d(
    Real* output_Y,int out_n, int out_c, int out_d, int out_h, int out_w,
    Real* input_W, int k_n, int k_c, int k_d, int k_h, int k_w,
    int p_d, int p_h, int p_w,
    int s_h, int s_w,
    int d_h, int d_w,
    Real* input_X, int in_n, int in_c, int in_d, int in_h, int in_w
   )
{

    int tmp_h_length = out_h + k_h - 1; 
    int tmp_w_length = out_w + k_w - 1;
    int tmp_d_length = out_d + k_d - 1;

    Real* tmp_input_X = (Real*) malloc(in_n * in_c * tmp_d_length * tmp_h_length * tmp_w_length * sizeof(Real));
    
    convfft_3d_4backdata(output_Y, out_n, out_c, out_d, out_h, out_w,
        input_W, k_n, k_c, k_d, k_h, k_w,
        p_d , p_h,  p_w,
         s_h,  s_w,
         d_h,  d_w,
         tmp_input_X, in_n, in_c, tmp_d_length, tmp_h_length, tmp_w_length);


    matrixTrimBackData(tmp_input_X, input_X, in_n, in_c, tmp_d_length, tmp_h_length, tmp_w_length, in_d, in_h, in_w);
    
}
void conv_back_data_fft_3d_v2(
    Real* output_Y,int out_n, int out_c, int out_d, int out_h, int out_w,
    Real* input_W, int k_n, int k_c, int k_d, int k_h, int k_w,
    int p_d, int p_h, int p_w,
     int s_h, int s_w,
     int d_h, int d_w,
    Real* input_X, int in_n, int in_c, int in_d, int in_h, int in_w
   )
{
    Real* input_W_trans = (Real*) malloc( k_n * k_c * k_d * k_h * k_w *sizeof(Real));
    transposeByBlock( input_W, input_W_trans , k_n, k_c, k_d,k_h, k_w);
  convfft_3d_v2(output_Y, out_n, out_c, out_d, out_h, out_w,
       input_W_trans , k_c, k_n, k_d, k_h, k_w,
        in_d - out_d + 1, in_h - out_h + 1, in_w - out_w + 1,
        s_h,  s_w,
        d_h,  d_w,
        input_X, in_n, in_c, in_d, in_h, in_w);
}
__global__ void transposeByBlock(Real* input, int n, int c, int d, int h, int w, Real* output)
{
  int tdx = blockDim.x * blockIdx.x + threadIdx.x;
  if (tdx < n * c * d * h * w)
  {
    
    int index_n = tdx / (c * d * h * w);
    int index_c = (tdx % (c * d * h * w)) / (d * h * w);
    int index_d = ((tdx % (c * d * h * w)) % (d * h * w)) / (h * w);
    int index_h = (((tdx % (c * d * h * w)) % (d * h * w)) % (h * w)) / w;
    int index_w = (((tdx % (c * d * h * w)) % (d * h * w)) % (h * w)) % w;

   output[index_n * c * d * h * w + index_c * d * h * w + index_d * h * w + index_h * w + index_w] = input[index_c * c * d * h * w + index_n * d * h * w + index_d * h * w + index_h * w + index_w];
  }
}

void conv_back_filter_fft_2d_gpu(
    Real* input_X, int in_n, int in_c, int in_d, int in_h, int in_w,
    Real* output_Y,int out_n, int out_c, int out_d, int out_h, int out_w,
    int p_h, int p_w,
    int s_h, int s_w,
    int d_h, int d_w,
    Real* input_W, int k_n, int k_c, int k_d, int k_h, int k_w)
{

    Real* input_X_trans  = (Real*) malloc(in_n * in_c * in_d * in_h * in_w * sizeof(Real));
    Real* output_Y_trans = (Real*) malloc(out_n * out_c * out_d * out_h * out_w * sizeof(Real));
    Real* input_W_trans  = (Real*) malloc( k_n * k_c * k_d * k_h * k_w*sizeof(Real));

    // invoke the kernel 
    transposeByBlock(input_X,input_X_trans, in_n, in_c, in_h, in_w);
    transposeByBlock(output_Y,output_Y_trans, out_n, out_c, out_h, out_w);
    rot180Degree(output_Y_trans,out_c, out_n, out_d, out_h,  out_w);
    convfft_2d_gpu_v7(input_X_trans, in_c, in_n, in_d, in_h, in_w,
        output_Y_trans, out_c,out_n,  out_d, out_h,  out_w,
         p_h,  p_w,
         s_h,  s_w,
         d_h,  d_w,
         input_W_trans,  k_c, k_n,k_d, k_h, k_w);

    transposeByBlock( input_W_trans ,input_W,  k_c, k_n,k_h, k_w);
   
    rot180Degree(input_W,k_n, k_c, k_d, k_h, k_w);

        free(input_X_trans);
        free(output_Y_trans);
        free(input_W_trans);
}


void conv_back_data_fft_2d_gpu(
    Real* output_Y,int out_n, int out_c, int out_d, int out_h, int out_w,
    Real* input_W, int k_n, int k_c, int k_d, int k_h, int k_w,
    int p_h, int p_w,
    int s_h, int s_w,
    int d_h, int d_w,
    Real* input_X, int in_n, int in_c, int in_d, int in_h, int in_w
   )
{
    Real* input_W_trans = (Real*) malloc( k_n * k_c * k_d * k_h * k_w*sizeof(Real));
    
    transposeByBlock( input_W, input_W_trans , k_n, k_c, k_h, k_w);
    
    rot180Degree(input_W_trans, k_c , k_n , k_d , k_h , k_w);

    convfft_2d_gpu_v7(output_Y, out_n, out_c, out_d, out_h,  out_w,
        input_W_trans , k_c, k_n, k_d, k_h, k_w,
         p_h,  p_w,
         s_h,  s_w,
         d_h,  d_w,
         input_X, in_n, in_c, in_d, in_h, in_w);
}


void conv_back_data_fft_3d_gpu(
    Real* output_Y,int out_n, int out_c, int out_d, int out_h, int out_w,
    Real* input_W, int k_n, int k_c, int k_d, int k_h, int k_w,
    int p_d, int p_h, int p_w,
    int s_d, int s_h, int s_w,
    int d_d, int d_h, int d_w,
    Real* input_X, int in_n, int in_c, int in_d, int in_h, int in_w
   )
{

    int tmp_h_length = 1 + (out_h  + 2 * 1 - ((k_h - 1) * 1 + 1))/1 ;
    int tmp_w_length = 1 + (out_w  + 2 * 1 - ((k_w - 1) * 1 + 1))/1 ;
    int tmp_d_length = (1 + (out_d  + 2 * 1 - ((k_d - 1) * 1 + 1))/1) < 0 ? in_d + 1 : 1 + (out_d  + 2 * 1 - ((k_d - 1) * 1 + 1))/1 ;    

    Real* tmp_input_X = (Real*) malloc(in_n * in_c * tmp_d_length * tmp_h_length * tmp_w_length * sizeof(Real));
    
    convfft_3d_gpu(output_Y, out_n, out_c, out_d, out_h, out_w,
        input_W, k_n, k_c, k_d, k_h, k_w,
        p_d , p_h,  p_w,
        s_d, s_h,  s_w,
        d_d, d_h,  d_w,
        tmp_input_X, in_n, in_c, tmp_d_length, tmp_h_length, tmp_w_length);


    matrixTrimBackData(tmp_input_X, input_X, in_n, in_c, tmp_d_length, tmp_h_length, tmp_w_length, in_d, in_h, in_w);
    
}

void conv_back_data_fft_3d_gpu_v2(
    Real* output_Y,int out_n, int out_c, int out_d, int out_h, int out_w,
    Real* input_W, int k_n, int k_c, int k_d, int k_h, int k_w,
    int p_d, int p_h, int p_w,
    int s_d, int s_h, int s_w,
    int d_d, int d_h, int d_w,
    Real* input_X, int in_n, int in_c, int in_d, int in_h, int in_w
   )
{

    Real* input_W_trans = (Real*) malloc( k_n * k_c * k_d * k_h * k_w *sizeof(Real));
    
    transposeByBlock( input_W, input_W_trans , k_n, k_c, k_d,k_h, k_w);
    
    convfft_3d_gpu(output_Y, out_n, out_c, out_d, out_h, out_w,
       input_W_trans , k_c, k_n, k_d, k_h, k_w,
        in_d - out_d + 1, in_h - out_h + 1, in_w - out_w + 1,
        s_d, s_h,  s_w,
        d_d, d_h,  d_w,
        input_X, in_n, in_c, in_d, in_h, in_w);
    
}

